//
//  BinaryAVLTree.swift
//  BinaryTree
//
//  Created by 郭红旗 on 2021/11/19.
//

import UIKit

class BinaryAVLTree<T:Equatable>:BinarySearchTree<T>{
    
    var balanceFator:Int = 0
    
    /// - Parameters: 重写父类方法
    override func createBinaryTree(element: T, parent: BinarySearchTree<T>.Node<T>?) -> BinarySearchTree<T>.Node<T> {
        return AVLNode<T>.init(element: element, parent: parent)
    }
    
    

    override func afterAdd(_ node: BinarySearchTree<T>.Node<T>) {
        var holderNode:Node<T>? = node
        while  holderNode?.parent != nil {
            holderNode = holderNode?.parent
            if isBalance(node: holderNode!) {
                // 更新高度
                updateHeight(node: holderNode!)
            }else{
                // 恢复平衡
                reBalance(grandNode: holderNode!)
                break
            }
        }
    }
    
    override func afterRemove(_ node: BinarySearchTree<T>.Node<T>) {
        var holderNode:Node<T>? = node
        while  holderNode?.parent != nil {
            holderNode = holderNode?.parent
            if isBalance(node: holderNode!) {
                // 更新高度
                updateHeight(node: holderNode!)
            }else{
                // 恢复平衡
                reBalance(grandNode: holderNode!)
            }
        }
    }
    
    
    /// 恢复平衡
    func reBalance(grandNode:Node<T>){
        let newNode = grandNode as! AVLNode<T>
        let parent = newNode.tallerChild()
        let node = parent.tallerChild()
        if parent.isLeftChild() {
            if node.isLeftChild() {
                // LL
                rotateRight(grand: newNode)
            }else{
                // LR
                rotateLeft(grand: parent)
                rotateRight(grand: newNode)
            }
        }else{
            if node.isLeftChild() {
                // RL
                rotateRight(grand: parent)
                rotateLeft(grand: newNode)
            }else{
                // RR
                rotateLeft(grand: newNode)
            }
        }

    }
    /// 左旋转
    func rotateLeft(grand:Node<T>){
        guard let parent = grand.right else {
            return
        }
        let child = parent.left
        grand.right = child
        parent.left = grand
        
        parent.parent = grand.parent
        if grand.isLeftChild() {
            grand.parent?.left = parent
        }else if(grand.isRightChild()){
            grand.parent?.right = parent
        }else{ // grand 是root节点
            self.root = parent
        }
        // 更新child和grand的parent
        grand.parent = parent
        child?.parent = grand
        
        // 更新高度
        updateHeight(node: grand)
        updateHeight(node: parent)
     }
    /// 右旋转
    func rotateRight(grand:Node<T>){
        guard let parent = grand.left else {
            return
        }
        let child = parent.right
        grand.left = child
        parent.right = grand
        parent.parent = grand.parent
        
        if grand.isLeftChild() {
            grand.parent?.left = parent
        }else if grand.isRightChild(){
            grand.parent?.right = parent
        }else{
            self.root = parent
        }
        
        grand.parent = parent
        child?.parent = grand

        // 更新高度
        updateHeight(node: grand)
        updateHeight(node: parent)
    }
    
    /// 高度获取
    private func nodeHeight(node:Node<T>?) -> Int{
        guard let new = node else {
            return 0
        }
        return 1 + max(nodeHeight(node: new.left), nodeHeight(node: new.right))
    }
    
    /// 更新高度
    private func updateHeight(node:Node<T>) {
        let newNode = node as! AVLNode<T>
        newNode.updateHeight()
    }
    
    /// 是否平衡 左右子节点绝对值小于1
    private func isBalance(node:Node<T>) -> Bool {
        let newNode = node as! AVLNode<T>
        return abs(newNode.banlanceFactor()) <= 1
    }
    
    
    class AVLNode<E:Equatable>:Node<E>{
        var height:Int = 1
        override init(element: E, parent: BinarySearchTree<T>.Node<E>?) {
            super.init(element: element, parent: parent)
        }
        /// 平衡因子
        func banlanceFactor() -> Int {
            var leftHeight = 0
            if let left = self.left as? AVLNode<E> {
                leftHeight = left.height
            }
            var rightHeight = 0
            if let right = self.right as? AVLNode<E> {
                rightHeight = right.height
            }
            return leftHeight - rightHeight
        }
        
        func updateHeight(){
            var leftHeight = 0
            if let left = self.left as? AVLNode<E> {
                leftHeight = left.height
            }
            var rightHeight = 0
            if let right = self.right as? AVLNode<E> {
                rightHeight = right.height
            }
            height = 1 + max(leftHeight, rightHeight)
        }
        
        func
        tallerChild() -> AVLNode<E>{
            var leftHeight = 0
            if let left = self.left as? AVLNode<E> {
                leftHeight = left.height
            }
            var rightHeight = 0
            if let right = self.right as? AVLNode<E> {
                rightHeight = right.height
            }
            if leftHeight > rightHeight {
                return self.left as! AVLNode
            }
            if rightHeight > leftHeight {
                return self.right as! AVLNode
            }
            return isLeftChild() ? self.left as! AVLNode : self.right as! AVLNode
        }

    }
    
    
}
